{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1ed03b3-b948-4af8-badf-d0c48491a8c8",
   "metadata": {},
   "source": [
    "# function_calling案例实战"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ff52aa6-556b-4c9c-b6b2-184d7d76cfde",
   "metadata": {},
   "source": [
    "## 1. 天气预报案例实战"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dc53a51-040e-40f1-9819-9dbe7c5043b7",
   "metadata": {},
   "source": [
    "### 1.1 准备环境"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a9ae778b-5077-4d71-b047-e25d4f6fd6ce",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: scipy in d:\\program files (x86)\\python3\\lib\\site-packages (1.12.0)\n",
      "Requirement already satisfied: tenacity in d:\\program files (x86)\\python3\\lib\\site-packages (8.2.3)\n",
      "Requirement already satisfied: tiktoken in d:\\program files (x86)\\python3\\lib\\site-packages (0.6.0)\n",
      "Collecting termcolor\n",
      "  Downloading termcolor-2.4.0-py3-none-any.whl (7.7 kB)\n",
      "Requirement already satisfied: openai in d:\\program files (x86)\\python3\\lib\\site-packages (1.14.3)\n",
      "Requirement already satisfied: requests in d:\\program files (x86)\\python3\\lib\\site-packages (2.31.0)\n",
      "Requirement already satisfied: numpy<1.29.0,>=1.22.4 in d:\\program files (x86)\\python3\\lib\\site-packages (from scipy) (1.26.4)\n",
      "Requirement already satisfied: regex>=2022.1.18 in d:\\program files (x86)\\python3\\lib\\site-packages (from tiktoken) (2023.12.25)\n",
      "Requirement already satisfied: sniffio in d:\\program files (x86)\\python3\\lib\\site-packages (from openai) (1.3.1)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in d:\\program files (x86)\\python3\\lib\\site-packages (from openai) (1.9.0)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in d:\\program files (x86)\\python3\\lib\\site-packages (from openai) (4.3.0)\n",
      "Requirement already satisfied: typing-extensions<5,>=4.7 in d:\\program files (x86)\\python3\\lib\\site-packages (from openai) (4.10.0)\n",
      "Requirement already satisfied: pydantic<3,>=1.9.0 in d:\\program files (x86)\\python3\\lib\\site-packages (from openai) (2.6.4)\n",
      "Requirement already satisfied: tqdm>4 in d:\\program files (x86)\\python3\\lib\\site-packages (from openai) (4.66.2)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in d:\\program files (x86)\\python3\\lib\\site-packages (from openai) (0.27.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in d:\\program files (x86)\\python3\\lib\\site-packages (from requests) (3.6)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in d:\\program files (x86)\\python3\\lib\\site-packages (from requests) (3.3.2)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in d:\\program files (x86)\\python3\\lib\\site-packages (from requests) (2.2.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in d:\\program files (x86)\\python3\\lib\\site-packages (from requests) (2024.2.2)\n",
      "Requirement already satisfied: exceptiongroup>=1.0.2 in d:\\program files (x86)\\python3\\lib\\site-packages (from anyio<5,>=3.5.0->openai) (1.2.0)\n",
      "Requirement already satisfied: httpcore==1.* in d:\\program files (x86)\\python3\\lib\\site-packages (from httpx<1,>=0.23.0->openai) (1.0.5)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in d:\\program files (x86)\\python3\\lib\\site-packages (from httpcore==1.*->httpx<1,>=0.23.0->openai) (0.14.0)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in d:\\program files (x86)\\python3\\lib\\site-packages (from pydantic<3,>=1.9.0->openai) (0.6.0)\n",
      "Requirement already satisfied: pydantic-core==2.16.3 in d:\\program files (x86)\\python3\\lib\\site-packages (from pydantic<3,>=1.9.0->openai) (2.16.3)\n",
      "Requirement already satisfied: colorama in d:\\program files (x86)\\python3\\lib\\site-packages (from tqdm>4->openai) (0.4.6)\n",
      "Installing collected packages: termcolor\n",
      "Successfully installed termcolor-2.4.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[notice] A new release of pip is available: 23.0.1 -> 24.0\n",
      "[notice] To update, run: python.exe -m pip install --upgrade pip\n"
     ]
    }
   ],
   "source": [
    "! pip install scipy tenacity tiktoken termcolor openai requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "4d8f2bcb-8998-4492-a687-6347fbcf5524",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from openai import OpenAI\n",
    "from tenacity import retry, wait_random_exponential, stop_after_attempt\n",
    "from termcolor import colored\n",
    "\n",
    "GPT_MODEL = \"gpt-3.5-turbo-0613\"\n",
    "client = OpenAI()\n",
    "\n",
    "#client.models.list()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b92a4df3-c955-48d8-8d39-130c21c2c03c",
   "metadata": {},
   "source": [
    "### 1.2 封装模型调用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "da788276-80cf-4b94-829e-92e1c69ab210",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 定义一个函数chat_completion_request，主要用于发送 聊天补全 请求到OpenAI服务器\n",
    "def chat_completion_request(messages, tools=None, tool_choice=None, model=GPT_MODEL):\n",
    "    try:\n",
    "        response = client.chat.completions.create(\n",
    "            model=model,\n",
    "            messages=messages,\n",
    "            tools=tools,\n",
    "            tool_choice=tool_choice,\n",
    "        )\n",
    "        return response\n",
    "    except Exception as e:\n",
    "        print(\"Unable to generate ChatCompletion response\")\n",
    "        print(f\"Exception: {e}\")\n",
    "        return e\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 1.3 定义打印函数"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "79665a4e-8016-46f1-afe2-cdc6fd02caf8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 定义一个函数pretty_print_conversation，用于打印消息对话内容\n",
    "def pretty_print_conversation(messages):\n",
    "\n",
    "    # 为不同角色设置不同的颜色\n",
    "    role_to_color = {\n",
    "        \"system\": \"red\",\n",
    "        \"user\": \"green\",\n",
    "        \"assistant\": \"blue\",\n",
    "        \"function\": \"magenta\",\n",
    "    }\n",
    "\n",
    "    # 遍历消息列表\n",
    "    for message in messages:\n",
    "            # 如果消息的角色是\"system\"，则用红色打印“content”\n",
    "            if isinstance(message, dict) and message[\"role\"] == \"system\":\n",
    "                print(colored(f\"system: {message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
    "            # 如果消息的角色是\"user\"，则用绿色打印“content”\n",
    "            elif isinstance(message, dict) and message[\"role\"] == \"user\":\n",
    "                print(colored(f\"user: {message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
    "            # 如果消息的角色是\"assistant\"，并且消息中包含\"function_call\"，则用蓝色打印\"function_call\"\n",
    "            elif hasattr(message, 'role') and message.role == \"assistant\" and message.function_call:\n",
    "                print(colored(f\"assistant: {message.function_call}\\n\", role_to_color[message.role]))\n",
    "            # 如果消息的角色是\"assistant\"，但是消息中不包含\"function_call\"，则用蓝色打印“content”\n",
    "            elif hasattr(message, 'role') and message.role == \"assistant\" and not message.function_call:\n",
    "                print(colored(f\"assistant: {message.content}\\n\", role_to_color[message.role]))\n",
    "            elif isinstance(message, dict) and message[\"role\"] == \"function\":\n",
    "            # 如果消息的角色是\"function\"，则用品红色打印“function”\n",
    "                print(colored(f\"function ({message['name']}): {message['content']}\\n\", role_to_color[message[\"role\"]]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 1.4 定义function函数"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "outputs": [],
   "source": [
    "# 定义一个名为functions的列表，其中包含两个字典，这两个字典分别定义了两个功能的相关参数\n",
    "\n",
    "# 第一个字典定义了一个名为\"get_current_weather\"的函数\n",
    "tools = [\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": \"get_current_weather\",# 函数的名称\n",
    "            \"description\": \"Get the current weather\",# 函数的描述\n",
    "            \"parameters\": {# 定义该函数需要的参数\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {# 参数的属性\n",
    "                    \"location\": {# 地点参数\n",
    "                        \"type\": \"string\",# 参数类型为字符串\n",
    "                        \"description\": \"The city and state, e.g. San Francisco, CA\",# 参数的描述\n",
    "                    },\n",
    "                    \"format\": {# 温度单位参数\n",
    "                        \"type\": \"string\",# 参数类型为字符串\n",
    "                        \"enum\": [\"celsius\", \"fahrenheit\"],# 摄氏度，华氏 参数的取值\n",
    "                        \"description\": \"The temperature unit to use. Infer this from the users location.\",# 参数的描述\n",
    "                    },\n",
    "                },\n",
    "                \"required\": [\"location\", \"format\"],# 该函数需要的必要参数\n",
    "            },\n",
    "        }\n",
    "    },\n",
    "    # 第二个字典定义了一个名为\"get_n_day_weather_forecast\"的函数\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": \"get_n_day_weather_forecast\",\n",
    "            \"description\": \"Get an N-day weather forecast\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"location\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The city and state, e.g. San Francisco, CA\",\n",
    "                    },\n",
    "                    \"format\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"enum\": [\"celsius\", \"fahrenheit\"],\n",
    "                        \"description\": \"The temperature unit to use. Infer this from the users location.\",\n",
    "                    },\n",
    "                    \"num_days\": {\n",
    "                        \"type\": \"integer\",\n",
    "                        \"description\": \"The number of days to forecast\",\n",
    "                    }\n",
    "                },\n",
    "                \"required\": [\"location\", \"format\", \"num_days\"]\n",
    "            },\n",
    "        }\n",
    "    },\n",
    "]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 1.5 函数调用：咨询今天天气"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "27789b41-fbcd-40ad-b8e2-8acf1507c453",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001B[31msystem: Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\n",
      "\u001B[0m\n",
      "\u001B[32muser: What's the weather like today\n",
      "\u001B[0m\n",
      "\u001B[34massistant: Sure, could you please provide the location for which you would like to know the weather?\n",
      "\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "messages = []\n",
    "messages.append({\"role\": \"system\", \"content\": \"Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\"})\n",
    "messages.append({\"role\": \"user\", \"content\": \"What's the weather like today\"})\n",
    "chat_response = chat_completion_request(\n",
    "    messages, tools=tools\n",
    ")\n",
    "assistant_message = chat_response.choices[0].message\n",
    "messages.append(assistant_message)\n",
    "#assistant_message\n",
    "# 打印助手的回复消息\n",
    "pretty_print_conversation(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 1.6 函数调用：回复当前地点"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001B[31msystem: Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\n",
      "\u001B[0m\n",
      "\u001B[32muser: What's the weather like today\n",
      "\u001B[0m\n",
      "\u001B[34massistant: Sure, could you please provide the location for which you would like to know the weather?\n",
      "\u001B[0m\n",
      "\u001B[32muser: I'm in ShenZhen, China.\n",
      "\u001B[0m\n",
      "\u001B[34massistant: None\n",
      "\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "messages.append({\"role\": \"user\", \"content\": \"I'm in ShenZhen, China.\"})\n",
    "#大模型自己选择是否使用函数\n",
    "#chat_response = chat_completion_request(messages, tools=tools)\n",
    "#强制使用函数\n",
    "chat_response = chat_completion_request(messages, tools=tools, tool_choice={\"type\": \"function\", \"function\": {\"name\": \"get_current_weather\"}})\n",
    "assistant_message = chat_response.choices[0].message\n",
    "messages.append(assistant_message)\n",
    "#assistant_message\n",
    "# 打印助手的回复消息\n",
    "pretty_print_conversation(messages)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 设置问答样例，用于咨询未来5天的天气"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001B[31msystem: Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\n",
      "\u001B[0m\n",
      "\u001B[32muser: what is the weather going to be like in Glasgow, Scotland over the next x days\n",
      "\u001B[0m\n",
      "\u001B[34massistant: Sure, I can help you with that. Could you please specify the value of `x`?\n",
      "\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "messages = []\n",
    "messages.append({\"role\": \"system\", \"content\": \"Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\"})\n",
    "messages.append({\"role\": \"user\", \"content\": \"what is the weather going to be like in Glasgow, Scotland over the next x days\"})\n",
    "chat_response = chat_completion_request(\n",
    "    messages, tools=tools\n",
    ")\n",
    "assistant_message = chat_response.choices[0].message\n",
    "messages.append(assistant_message)\n",
    "#assistant_message\n",
    "# 打印助手的回复消息\n",
    "pretty_print_conversation(messages)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "### 咨询未来5天天气"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "outputs": [
    {
     "data": {
      "text/plain": "Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, role='assistant', function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_ihsCi6R5mIRHQPhXPQC9DLex', function=Function(arguments='{\\n  \"location\": \"Glasgow, Scotland\",\\n  \"format\": \"celsius\",\\n  \"num_days\": 5\\n}', name='get_n_day_weather_forecast'), type='function')]))"
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages.append({\"role\": \"user\", \"content\": \"5 days\"})\n",
    "chat_response = chat_completion_request(\n",
    "    messages, tools=tools\n",
    ")\n",
    "chat_response.choices[0]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "### 强制使用函数"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "outputs": [
    {
     "data": {
      "text/plain": "ChatCompletionMessage(content=None, role='assistant', function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_5ll1eH1bGp2QvNziUUe5wnzD', function=Function(arguments='{\\n  \"location\": \"Toronto, Canada\",\\n  \"format\": \"celsius\",\\n  \"num_days\": 1\\n}', name='get_n_day_weather_forecast'), type='function')])"
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# in this cell we force the model to use get_n_day_weather_forecast\n",
    "messages = []\n",
    "messages.append({\"role\": \"system\", \"content\": \"Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\"})\n",
    "messages.append({\"role\": \"user\", \"content\": \"Give me a weather report for Toronto, Canada.\"})\n",
    "chat_response = chat_completion_request(\n",
    "    messages, tools=tools, tool_choice={\"type\": \"function\", \"function\": {\"name\": \"get_n_day_weather_forecast\"}}\n",
    ")\n",
    "chat_response.choices[0].message"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "### 强制不使用函数"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "outputs": [
    {
     "data": {
      "text/plain": "ChatCompletionMessage(content='{\\n  \"location\": \"Toronto, Canada\",\\n  \"format\": \"celsius\"\\n}', role='assistant', function_call=None, tool_calls=None)"
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = []\n",
    "messages.append({\"role\": \"system\", \"content\": \"Don't make assumptions about what values to plug into functions. Ask for clarification if a user request is ambiguous.\"})\n",
    "messages.append({\"role\": \"user\", \"content\": \"Give me the current weather (use Celcius) for Toronto, Canada.\"})\n",
    "chat_response = chat_completion_request(\n",
    "    messages, tools=tools, tool_choice=\"none\"\n",
    ")\n",
    "chat_response.choices[0].message"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 2. 基于大模型实现数据平台查询"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "id": "e3b858cf-1dcc-4fb0-9210-9fe66d49c55b",
   "metadata": {},
   "source": [
    "### 1. 定义数据库信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "6065e696-0796-4630-93f6-b73b810bc26a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Opened database successfully\n"
     ]
    }
   ],
   "source": [
    "import sqlite3\n",
    "\n",
    "conn = sqlite3.connect(\"data/chinook.db\")\n",
    "print(\"Opened database successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32519aac-6a53-48a2-af9c-c2b9c6a76911",
   "metadata": {},
   "source": [
    "首先定义三个函数get_table_names、get_column_names和get_database_info，用于从数据库连接对象中获取数据库的表名、表的列名以及整体数据库的信息。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "1e534820-200e-4789-8fc6-5de8e0820f54",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_table_names(conn):\n",
    "    \"\"\"返回一个包含所有表名的列表\"\"\"\n",
    "    table_names = []  # 创建一个空的表名列表\n",
    "    # 执行SQL查询，获取数据库中所有表的名字\n",
    "    tables = conn.execute(\"SELECT name FROM sqlite_master WHERE type='table';\")\n",
    "    # 遍历查询结果，并将每个表名添加到列表中\n",
    "    for table in tables.fetchall():\n",
    "        table_names.append(table[0])\n",
    "    return table_names  # 返回表名列表\n",
    "\n",
    "\n",
    "def get_column_names(conn, table_name):\n",
    "    \"\"\"返回一个给定表的所有列名的列表\"\"\"\n",
    "    column_names = []  # 创建一个空的列名列表\n",
    "    # 执行SQL查询，获取表的所有列的信息\n",
    "    columns = conn.execute(f\"PRAGMA table_info('{table_name}');\").fetchall()\n",
    "    # 遍历查询结果，并将每个列名添加到列表中\n",
    "    for col in columns:\n",
    "        column_names.append(col[1])\n",
    "    return column_names  # 返回列名列表\n",
    "\n",
    "\n",
    "def get_database_info(conn):\n",
    "    \"\"\"返回一个字典列表，每个字典包含一个表的名字和列信息\"\"\"\n",
    "    table_dicts = []  # 创建一个空的字典列表\n",
    "    # 遍历数据库中的所有表\n",
    "    for table_name in get_table_names(conn):\n",
    "        columns_names = get_column_names(conn, table_name)  # 获取当前表的所有列名\n",
    "        # 将表名和列名信息作为一个字典添加到列表中\n",
    "        table_dicts.append({\"table_name\": table_name, \"column_names\": columns_names})\n",
    "    return table_dicts  # 返回字典列表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "ff5fe286-57f0-4c9e-a2a9-85bbdacb36ba",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "'Table: albums\\nColumns: AlbumId, Title, ArtistId\\nTable: sqlite_sequence\\nColumns: name, seq\\nTable: artists\\nColumns: ArtistId, Name\\nTable: customers\\nColumns: CustomerId, FirstName, LastName, Company, Address, City, State, Country, PostalCode, Phone, Fax, Email, SupportRepId\\nTable: employees\\nColumns: EmployeeId, LastName, FirstName, Title, ReportsTo, BirthDate, HireDate, Address, City, State, Country, PostalCode, Phone, Fax, Email\\nTable: genres\\nColumns: GenreId, Name\\nTable: invoices\\nColumns: InvoiceId, CustomerId, InvoiceDate, BillingAddress, BillingCity, BillingState, BillingCountry, BillingPostalCode, Total\\nTable: invoice_items\\nColumns: InvoiceLineId, InvoiceId, TrackId, UnitPrice, Quantity\\nTable: media_types\\nColumns: MediaTypeId, Name\\nTable: playlists\\nColumns: PlaylistId, Name\\nTable: playlist_track\\nColumns: PlaylistId, TrackId\\nTable: tracks\\nColumns: TrackId, Name, AlbumId, MediaTypeId, GenreId, Composer, Milliseconds, Bytes, UnitPrice\\nTable: sqlite_stat1\\nColumns: tbl, idx, stat'"
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 获取数据库信息，并存储为字典列表\n",
    "database_schema_dict = get_database_info(conn)\n",
    "\n",
    "# 将数据库信息转换为字符串格式，方便后续使用\n",
    "database_schema_string = \"\\n\".join(\n",
    "    [\n",
    "        f\"Table: {table['table_name']}\\nColumns: {', '.join(table['column_names'])}\"\n",
    "        for table in database_schema_dict\n",
    "    ]\n",
    ")\n",
    "database_schema_string"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c47b12e-25e1-4ba5-89d6-0c77863aa579",
   "metadata": {},
   "source": [
    "### 2.定义 functions让 GPT 模型帮我们构造一个完整的 SQL 查询"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "bb495873-eb37-4d4f-82c7-1b6aefcc21d9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 定义一个功能列表，其中包含一个功能字典，该字典定义了一个名为\"ask_database_get_sql\"的功能，用于回答用户关于音乐的问题\n",
    "tools = [\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": \"ask_database_get_sql\",\n",
    "            \"description\": \"Use this function to answer user questions about music. Input should be a fully formed SQL query.\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"query\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": f\"\"\"\n",
    "                                SQL query extracting info to answer the user's question.\n",
    "                                SQL should be written using this database schema:\n",
    "                                {database_schema_string}\n",
    "                                The query should be returned in plain text, not in JSON.\n",
    "                                \"\"\",\n",
    "                    }\n",
    "                },\n",
    "                \"required\": [\"query\"],\n",
    "            },\n",
    "        }\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f9323ca-4bdf-4c9c-ad75-5d1d794e96bd",
   "metadata": {},
   "source": [
    "### 3. 执行SQL语句（根据大模型返回的SQL，执行查询操作）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "91054286-6cfa-4b4e-9ce2-1f67b3c14ceb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def ask_database(conn, query):\n",
    "    \"\"\"使用 query 来查询 SQLite 数据库的函数。\"\"\"\n",
    "    try:\n",
    "        results = str(conn.execute(query).fetchall())  # 执行查询，并将结果转换为字符串\n",
    "    except Exception as e:  # 如果查询失败，捕获异常并返回错误信息\n",
    "        results = f\"query failed with error: {e}\"\n",
    "    return results  # 返回查询结果\n",
    "\n",
    "\n",
    "def execute_function_call(message):\n",
    "    \"\"\"执行函数调用\"\"\"\n",
    "    # 判断功能调用的名称是否为 \"ask_database_get_sql\"\n",
    "    if message.tool_calls[0].function.name == \"ask_database_get_sql\":\n",
    "        # 如果是，则获取功能调用的参数，这里是 SQL 查询\n",
    "        query = json.loads(message.tool_calls[0].function.arguments)[\"query\"]\n",
    "        # 使用 ask_database 函数执行查询，并获取结果\n",
    "        results = ask_database(conn, query)\n",
    "    else:\n",
    "        # 如果功能调用的名称不是 \"ask_database\"，则返回错误信息\n",
    "        results = f\"Error: function {message.tool_calls[0].function.name} does not exist\"\n",
    "    return results  # 返回结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "f6eb4e24-759b-4f45-a728-df0ab757a917",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatCompletionMessage(content=None, role='assistant', function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_4lto09jLf2FSDjyTUYPTJU51', function=Function(arguments='{\\n  \"query\": \"SELECT artists.Name, COUNT(tracks.TrackId) as TrackCount FROM artists JOIN albums ON artists.ArtistId = albums.ArtistId JOIN tracks ON albums.AlbumId = tracks.AlbumId GROUP BY artists.ArtistId ORDER BY TrackCount DESC LIMIT 5\"\\n}', name='ask_database_get_sql'), type='function')])\n",
      "[('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Deep Purple', 92)]\n",
      "\u001B[31msystem: Answer user questions by generating SQL queries against the Chinook Music Database.\n",
      "\u001B[0m\n",
      "\u001B[32muser: Hi, who are the top 5 artists by number of tracks?\n",
      "\u001B[0m\n",
      "\u001B[35mfunction (ask_database_get_sql): [('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Deep Purple', 92)]\n",
      "\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "# 创建一个空的消息列表\n",
    "messages = []\n",
    "\n",
    "# 向消息列表中添加一个系统角色的消息，内容是 \"Answer user questions by generating SQL queries against the Chinook Music Database.\"\n",
    "messages.append({\"role\": \"system\", \"content\": \"Answer user questions by generating SQL queries against the Chinook Music Database.\"})\n",
    "\n",
    "# 向消息列表中添加一个用户角色的消息，内容是 \"Hi, who are the top 5 artists by number of tracks?\"\n",
    "## 根据曲目数量排名前五的艺术家是谁？\n",
    "messages.append({\"role\": \"user\", \"content\": \"Hi, who are the top 5 artists by number of tracks?\"})\n",
    "\n",
    "# 使用 chat_completion_request 函数获取聊天响应\n",
    "chat_response = chat_completion_request(messages, tools)\n",
    "\n",
    "# 从聊天响应中获取助手的消息\n",
    "assistant_message = chat_response.choices[0].message\n",
    "#查看大模型，调用自定义函数生成的SQL\n",
    "#print(assistant_message)\n",
    "\n",
    "# 将助手的消息添加到消息列表中\n",
    "assistant_message.content = str(assistant_message.tool_calls[0].function)\n",
    "messages.append({\"role\": assistant_message.role, \"content\": assistant_message.content})\n",
    "\n",
    "# 如果助手的消息中有功能调用\n",
    "if assistant_message.tool_calls:\n",
    "    # 使用 execute_function_call 函数执行功能调用，并获取结果\n",
    "    results = execute_function_call(assistant_message)\n",
    "    #print(results)\n",
    "    # 将功能的结果作为一个功能角色的消息添加到消息列表中\n",
    "    messages.append({\"role\": \"function\", \"tool_call_id\": assistant_message.tool_calls[0].id, \"name\": assistant_message.tool_calls[0].function.name, \"content\": results})\n",
    "\n",
    "# 使用 pretty_print_conversation 函数打印对话\n",
    "pretty_print_conversation(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "37219a3d-0954-49dd-bbfa-482014cb3e6c",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "[{'table_name': 'albums', 'column_names': ['AlbumId', 'Title', 'ArtistId']},\n {'table_name': 'sqlite_sequence', 'column_names': ['name', 'seq']},\n {'table_name': 'artists', 'column_names': ['ArtistId', 'Name']},\n {'table_name': 'customers',\n  'column_names': ['CustomerId',\n   'FirstName',\n   'LastName',\n   'Company',\n   'Address',\n   'City',\n   'State',\n   'Country',\n   'PostalCode',\n   'Phone',\n   'Fax',\n   'Email',\n   'SupportRepId']},\n {'table_name': 'employees',\n  'column_names': ['EmployeeId',\n   'LastName',\n   'FirstName',\n   'Title',\n   'ReportsTo',\n   'BirthDate',\n   'HireDate',\n   'Address',\n   'City',\n   'State',\n   'Country',\n   'PostalCode',\n   'Phone',\n   'Fax',\n   'Email']},\n {'table_name': 'genres', 'column_names': ['GenreId', 'Name']},\n {'table_name': 'invoices',\n  'column_names': ['InvoiceId',\n   'CustomerId',\n   'InvoiceDate',\n   'BillingAddress',\n   'BillingCity',\n   'BillingState',\n   'BillingCountry',\n   'BillingPostalCode',\n   'Total']},\n {'table_name': 'invoice_items',\n  'column_names': ['InvoiceLineId',\n   'InvoiceId',\n   'TrackId',\n   'UnitPrice',\n   'Quantity']},\n {'table_name': 'media_types', 'column_names': ['MediaTypeId', 'Name']},\n {'table_name': 'playlists', 'column_names': ['PlaylistId', 'Name']},\n {'table_name': 'playlist_track', 'column_names': ['PlaylistId', 'TrackId']},\n {'table_name': 'tracks',\n  'column_names': ['TrackId',\n   'Name',\n   'AlbumId',\n   'MediaTypeId',\n   'GenreId',\n   'Composer',\n   'Milliseconds',\n   'Bytes',\n   'UnitPrice']},\n {'table_name': 'sqlite_stat1', 'column_names': ['tbl', 'idx', 'stat']}]"
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "database_schema_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "12dae1b1-a349-482a-bc02-dcf8abd1fd51",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "我是functions参数\n",
      "\u001B[31msystem: Answer user questions by generating SQL queries against the Chinook Music Database.\n",
      "\u001B[0m\n",
      "\u001B[32muser: Hi, who are the top 5 artists by number of tracks?\n",
      "\u001B[0m\n",
      "\u001B[34massistant[function_call]: {\n",
      "  \"name\": \"ask_database\",\n",
      "  \"arguments\": \"{\\n  \\\"query\\\": \\\"SELECT artists.Name, COUNT(tracks.TrackId) as NumTracks FROM artists JOIN albums ON artists.ArtistId = albums.ArtistId JOIN tracks ON albums.AlbumId = tracks.AlbumId GROUP BY artists.ArtistId ORDER BY NumTracks DESC LIMIT 5\\\"\\n}\"\n",
      "}\n",
      "\u001B[0m\n",
      "\u001B[35mfunction (ask_database): [('Iron Maiden', 213), ('U2', 135), ('Led Zeppelin', 114), ('Metallica', 112), ('Deep Purple', 92)]\n",
      "\u001B[0m\n",
      "\u001B[32muser: What is the name of the album with the most tracks?\n",
      "\u001B[0m\n",
      "\u001B[34massistant[function_call]: {\n",
      "  \"name\": \"ask_database\",\n",
      "  \"arguments\": \"{\\n  \\\"query\\\": \\\"SELECT albums.Title, COUNT(tracks.TrackId) as NumTracks FROM albums JOIN tracks ON albums.AlbumId = tracks.AlbumId GROUP BY albums.AlbumId ORDER BY NumTracks DESC LIMIT 1\\\"\\n}\"\n",
      "}\n",
      "\u001B[0m\n",
      "\u001B[35mfunction (ask_database): [('Greatest Hits', 57)]\n",
      "\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "# 向消息列表中添加一个用户的问题，内容是 \"What is the name of the album with the most tracks?\"\n",
    "## 哪张专辑的曲目数量最多？\n",
    "messages.append({\"role\": \"user\", \"content\": \"What is the name of the album with the most tracks?\"})\n",
    "\n",
    "# 使用 chat_completion_request 函数获取聊天响应\n",
    "chat_response = chat_completion_request(messages, tools)\n",
    "\n",
    "# 从聊天响应中获取助手的消息\n",
    "#assistant_message = chat_response[\"choices\"][0][\"message\"]\n",
    "assistant_message = chat_response.choices[0].message\n",
    "\n",
    "# 将助手的消息添加到消息列表中\n",
    "#messages.append(assistant_message)\n",
    "assistant_message.content = str(assistant_message.tool_calls[0].function)\n",
    "messages.append({\"role\": assistant_message.role, \"content\": assistant_message.content})\n",
    "\n",
    "# 如果助手的消息中有功能调用(即有解析返回SQL)\n",
    "if assistant_message.tool_calls:\n",
    "    # 根据大模型返回的SQL，使用 execute_function_call 函数执行功能调用，并获取结果\n",
    "    results = execute_function_call(assistant_message)\n",
    "    # 将功能的结果作为一个功能角色的消息添加到消息列表中\n",
    "    #messages.append({\"role\": \"function\", \"content\": results, \"name\": assistant_message[\"function_call\"][\"name\"]})\n",
    "    messages.append({\"role\": \"function\", \"tool_call_id\": assistant_message.tool_calls[0].id, \"name\": assistant_message.tool_calls[0].function.name, \"content\": results})\n",
    "\n",
    "# 使用 pretty_print_conversation 函数打印对话\n",
    "pretty_print_conversation(messages)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
